#ifndef BTL_ROBUSTESTIMATION_BASICRANSAC
#define BTL_ROBUSTESTIMATION_BASICRANSAC

#include "RANSAC.hpp"
#include <btl/Maths/RandomSampler.hpp>

#include <boost/optional.hpp>
#include <boost/random/random_number_generator.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <vector>

namespace btl {
namespace robustestimation {

template <class ModelType, int SampleSize, class RNG = boost::random_number_generator<boost::mt19937> >
class BasicRANSAC : public RobustModelFinder<ModelType, RNG>
{
   public:
      BasicRANSAC(double inlierThreshold, int numHypotheses, int minInliers, int refinementSteps);

      virtual boost::optional<ModelType> findModel(const ModelSearchSpace<ModelType>& searchSpace, RNG& rng);

      const std::vector<int>& getInliers() const;

   private:
      std::vector<int> _inliers;
      double _inlierThreshold;
      int _numHypotheses;
      int _minInliers;
      int _refinementSteps;
};

} //robustestimation
} //btl

namespace btl
{

using robustestimation::BasicRANSAC;

} //btl

// ====================================================================
// === Implementation

namespace btl {
namespace robustestimation {

template <class ModelType, int SampleSize, class RNG>
inline BasicRANSAC<ModelType,SampleSize,RNG>::BasicRANSAC(double inlierThreshold, int numHypotheses, int minInliers, int refinementSteps):
   _inlierThreshold(inlierThreshold), _numHypotheses(numHypotheses), _minInliers(minInliers), _refinementSteps(refinementSteps)
{
}

template <class ModelType, int SampleSize, class RNG>
inline const std::vector<int>& BasicRANSAC<ModelType,SampleSize,RNG>::getInliers() const
{
   return _inliers;
}

template <class ModelType, int SampleSize, class RNG>
inline boost::optional<ModelType> BasicRANSAC<ModelType,SampleSize,RNG>::findModel(
      const ModelSearchSpace<ModelType>& searchSpace, RNG& rng)
{
   const int N = searchSpace.numObservations();

   _inliers.clear();

   boost::optional<ModelType> bestModel;
   double bestCost = std::numeric_limits<double>::infinity();

   // observation indexes
   // (is shuffled to select samples)
   std::vector<int> obsIndexes;
   obsIndexes.reserve(N);
   for (int i = 0; i < N; ++i)
      obsIndexes.push_back(i);

   _inliers.clear();
   std::vector<int> inliers;
   inliers.reserve(N);

   for (int i = 0; i < _numHypotheses; ++i)
   {
      btl::shuffleMofN(
         obsIndexes.begin(),
         obsIndexes.begin() + SampleSize,
         obsIndexes.end(),
         rng);

      ModelType model = searchSpace.buildModel(&obsIndexes[0], SampleSize);

      // FIXME: use more numerically stable sum??
      //  probably won't make a big difference here because we're using doubles
      //  and we probably won't have enough observations to get major inaccuracies
      //  (I hope)
      inliers.clear();
      for (int i = 0; i < N; ++i)
      {
         double err = searchSpace.errorMetric(model, i);
         if (err < _inlierThreshold)
            inliers.push_back(i);
      }

      if (inliers.size() > _minInliers)
      {
         ModelType refined = searchSpace.buildModel(&inliers[0], inliers.size());

         double cost = 0.0;
         for (int i = 0; i < inliers.size(); ++i)
            cost += searchSpace.errorMetric(refined, inliers[i]);

         if (cost < bestCost)
         {
            bestCost = cost;
            bestModel = refined;
            _inliers.swap(inliers);
         }
      }
   }

   if (bestModel && _refinementSteps)
   {
      for (int step = 0; step < _refinementSteps; ++step)
      {
         _inliers.clear();
         for (int i = 0; i < N; ++i)
         {
            double err = searchSpace.errorMetric(*bestModel, i);
            if (err < _inlierThreshold)
               _inliers.push_back(i);
         }

         bestModel = searchSpace.buildModel(&_inliers[0], _inliers.size());
      }
   }

   return bestModel;
}

} //robustestimation
} //btl

#endif //BTL_ROBUSTESTIMATION_BASICRANSAC
